package io.gatling.extensions.mongodb.builder

import com.typesafe.scalalogging.StrictLogging
import io.gatling.core.action.Action
import io.gatling.core.action.builder.ActionBuilder
import io.gatling.core.check.Check
import io.gatling.core.session.Expression
import io.gatling.core.structure.ScenarioContext
import io.gatling.extensions.mongodb.action.MongoDBAction
import io.gatling.extensions.mongodb.protocol.MongoDBProtocol
import org.mongodb.scala.Document
import org.mongodb.scala.bson.conversions.Bson
import org.mongodb.scala.result._

case class MongoDBActionBuilder[R](name: String, dbName: String, collectionName: String) extends ActionBuilder with StrictLogging {


  var mongoDBOperate: MongoDBOperate = _

  def find(filterJson: Expression[String]): MongoDBActionBuilder[R] = {
    mongoDBOperate = FindOperate(dbName, collectionName, filterJson)
    type R = Document
    this
  }

  def insertOne(documentJson: Expression[String]): MongoDBActionBuilder[R] = {
    mongoDBOperate = InsertOneOperate(dbName, collectionName, documentJson)
    type R = InsertOneResult
    this
  }

  /**
    * 这里insert的内容不支持Expression
    */
  def insertMany(documents: Seq[_ <: Document]): MongoDBActionBuilder[R] = {
    mongoDBOperate = InsertManyOperate(dbName, collectionName, documents)
    type R = InsertManyResult
    this
  }

  def updateOne(filterJson: Expression[String], updateJson: Expression[String]): MongoDBActionBuilder[R] = {
    mongoDBOperate = UpdateOneOperate(dbName, collectionName, filterJson, updateJson)
    type R = UpdateResult
    this
  }

  /**
    * 这里update的内容不支持Expression
    */
  def updateMany(filterJson: Expression[String], update: Seq[Bson]): MongoDBActionBuilder[R] = {
    mongoDBOperate = UpdateManyOperate(dbName, collectionName, filterJson, update)
    type R = UpdateResult
    this
  }

  def deleteOne(filterJson: Expression[String]): MongoDBActionBuilder[R] = {
    mongoDBOperate = DeleteOneOperate(dbName, collectionName, filterJson)
    type R = DeleteResult
    this
  }

  def deleteMany(filterJson: Expression[String]): MongoDBActionBuilder[R] = {
    mongoDBOperate = DeleteManyOperate(dbName, collectionName, filterJson)
    type R = DeleteResult
    this
  }

  var checkList: List[Check[R]] = List.empty[Check[R]]

  def check(checks: Check[R]*): MongoDBActionBuilder[R] = {
    checkList = checkList ::: checks.toList
    this
  }

  override def build(ctx: ScenarioContext, next: Action): Action = {
    import ctx._
    MongoDBAction(
      name,
      mongoDBOperate,
      protocolComponentsRegistry.components(MongoDBProtocol.mongoDBProtocolKey),
      coreComponents,
      throttled,
      next,
      checkList
    )
  }
}

sealed trait MongoDBOperate

case class FindOperate(
                        dbName: String,
                        collectionName: String,
                        filter: Expression[String]
                      ) extends MongoDBOperate

case class InsertOneOperate(
                             dbName: String,
                             collectionName: String,
                             document: Expression[String]
                           ) extends MongoDBOperate

case class InsertManyOperate(
                              dbName: String,
                              collectionName: String,
                              documents: Seq[_ <: Document]
                            ) extends MongoDBOperate

case class UpdateOneOperate(
                             dbName: String,
                             collectionName: String,
                             filter: Expression[String],
                             update: Expression[String]
                           ) extends MongoDBOperate

case class UpdateManyOperate(
                              dbName: String,
                              collectionName: String,
                              filter: Expression[String],
                              update: Seq[Bson]
                            ) extends MongoDBOperate

case class DeleteOneOperate(
                             dbName: String,
                             collectionName: String,
                             filter: Expression[String]
                           ) extends MongoDBOperate

case class DeleteManyOperate(
                              dbName: String,
                              collectionName: String,
                              filter: Expression[String]
                            ) extends MongoDBOperate
